# from paraFit import *
from emulation import *
from PSOSearch import *
from GeneticSearch import genetic_altorithm
import os
import argparse

def str_to_bool(value):
    if value.lower() in ('True', 'true', '1', 't', 'y', 'yes'):
        return True
    elif value.lower() in ('False', 'false', '0', 'f', 'n', 'no'):
        return False
    else:
        raise argparse.ArgumentTypeError('Boolean value expected')

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument('--model', type=str, help="model name", required=True)
    parser.add_argument('--pref_batch_lower', type=int, help="pref_batch_lower", required=False, default=20)
    parser.add_argument('--pref_batch_upper', type=int, help="pref_batch_upper", required=False, default=50)
    parser.add_argument('--decode_batch_lower', type=int, help="decode_batch_lower", required=False, default=30)
    parser.add_argument('--decode_batch_upper', type=int, help="decode_batch_upper", required=False, default=150)
    parser.add_argument('--prefill_token_delay_tolerance_lower', type=int, help="prefill_token_delay_tolerance_lower", required=False, default=500)
    parser.add_argument('--prefill_token_delay_tolerance_upper', type=int, help="prefill_token_delay_tolerance_upper", required=False, default=800)
    parser.add_argument('--decode_constrains', type=float, help="decode_constrains", required=False, default=50)
    parser.add_argument('--is_firsttoken_constrained', type=str_to_bool, help="is_firsttoken_constrained", required=False, default=False)
    parser.add_argument('--firsttoken_constrains', type=float, help="firsttoken_constrains", required=False, default=1000)
    parser.add_argument('--num_populations', type=int, help="num_populations", required=False, default=4)
    parser.add_argument('--num_iterations', type=int, help="num_iterations", required=False, default=4)
    parser.add_argument('--concurrency', type=int, help="Concurrency", required=False, default=1000)
    parser.add_argument('--request_rate', type=int, help="Requestrate", required=False, default=20)
    parser.add_argument('--is_requestrate_compute', type=str_to_bool, help="Requestrate", required=False, default=False)
    parser.add_argument('--request_rate_lower', type=float, help="request_rate_lower", required=False, default=15)
    parser.add_argument('--request_rate_upper', type=float, help="request_rate_upper", required=False, default=25)
    parser.add_argument('--using_genetic', type=str_to_bool, help="Genetic Algorithm", required=False, default=False)
    parser.add_argument('--gene_length', type=int, help="gene_length", required=False, default=10)
    parser.add_argument('--is_SLO', type=str_to_bool, help="SLO", required=False, default=False)
    parser.add_argument('--is_P90', type=str_to_bool, help="P90", required=False, default=False)
    parser.add_argument('--no_constrain', type=str_to_bool, help="no_constrain", required=False, default=False)
    parser.add_argument('--is_prefixcache', type=str_to_bool, help="no_constrain", required=False, default=False)
    parser.add_argument('--is_splitfuse', type=str_to_bool, help="no_constrain", required=False, default=False)
    parser.add_argument('--data_name', type=str, help="no_constrain", required=False, default="short")
    parser.add_argument('--output_len', type=int, help="no_constrain", required=False, default=512)
    parser.add_argument('--supportSelectBatch', type=str_to_bool, help="supportSelectBatch", required=False, default=False)
    parser.add_argument('--is_speculative', type=str_to_bool, help="is_speculative", required=False, default=False)

    args = parser.parse_args()

    if args.is_requestrate_compute:
        bounds, hard_constrains, num_particles, gene_length = get_genetic_input2(args)
    else:
        bounds, hard_constrains, num_particles, gene_length = get_genetic_input1(args)

    print(f"模型名称: {args.model}")
    print("搜索空间：", bounds)
    print("粒子数量：", num_particles)
    if args.using_genetic:
        # bounds, hard_constrains, num_particles, gene_length = get_genetic_input1(args)

        print("基因长度：", gene_length)
        extra_param = {
            'config_file' : 'config/config.ini',
            'fitting_para_file' : 'config/config_llama3-8B.ini',
            'dataset_file' : 'config/config_dataset.ini',
            'model_file' : 'config/llama3-8b.ini',
            'maxPrefillBatchSize' : None,
            'maxPrefillTokens' : None
        }
        best_solution, best_fitness, decode_latency, prefill_latency = genetic_altorithm(
            genetic_algorithm_back_end, bounds, hard_constrains, num_particles, gene_length, extra_param, args.num_iterations, model_name = args.model, 
            concurrency = args.concurrency, request_rate = args.request_rate, isprefill_constrained= args.is_firsttoken_constrained
        )
    else:
        # bounds, hard_constrains, num_particles, iteration_num = get_genetic_input1(args)
        best_solution, best_fitness, decode_latency, prefill_latency = particle_swarm_optimization(
            particle_swarm_back_end, bounds, hard_constrains, num_particles, args.num_iterations, args
        )

    print("最优解：", best_solution)
    print("最优吞吐：", best_fitness)
    print("对应Decode时延：", decode_latency)
    print("对应首Token时延：", prefill_latency)